#include "poissondisk.h"
#include "../core/mathutil.h"
#include "../core/sampling.h"
#include "../core/geometry.h"

PDSampler::PDSampler(int xstart, int xend, int ystart, int yend, int spp, bool lns, bool ts) :
		Sampler(xstart, xend, ystart, yend, spp, lns, ts) {
	
	xPos = xstart;
	yPos = ystart;

	nPixelSamples = spp;	// aproximate, not actual number of samples

	samplePos = nPixelSamples;

	rng = new RNG(xstart + ystart *(xend - xstart));

	// Initialize data for Poisson disk sampling
	
	// is it a good estimation?
	radius = sqrtf(2./(nPixelSamples*M_PI));
	
	image.clear();
	lens.clear();
	time.clear();
	
}
PDSampler::~PDSampler() {
	SAFE_DELETE(rng);
}

Sampler *PDSampler::GetSubSampler(int num, int count) {
	int x0, x1, y0, y1;
	ComputeSubWindow(num, count, &x0, &x1, &y0, &y1);
	if (x0 == x1 || y0 == y1) 
		return NULL;
	return new PDSampler(x0, x1, y0, y1, nPixelSamples, needLensSamples, needTimeSamples);
}

int PDSampler::GetMoreSamples(Sample *sample) {

	if(samplePos == nPixelSamples) {

		if(xPos == xPixelEnd || yPos == yPixelEnd)
			return 0;

		// generate imageSample and lensSample using Poisson disk sampling
		bluenoiseSampling(radius, image, rng);
		if(needLensSamples) {
			bluenoiseSampling(radius, lens, rng);
			size_t n = max(image.size(), lens.size());
			for(uint i = image.size(); i < n; i++)
				image.push_back(Vector2(rng->RandomFloat(), rng->RandomFloat()));
			for(uint i = lens.size(); i < n; i++)
				lens.push_back(Vector2(rng->RandomFloat(), rng->RandomFloat()));
		}

		nPixelSamples = image.size();

		if(needTimeSamples)
			StratifiedSample1D(time, nPixelSamples, true, rng);

		if(needLensSamples)
			Shuffle(lens, nPixelSamples, rng);
		
		if(needTimeSamples)
			Shuffle(time, nPixelSamples, rng);
		
		samplePos = 0;
		// advance
		xPos++;
		if(xPos == xPixelEnd) {
			xPos = xPixelStart;
			yPos++;
		}
	}	
	
	// get samples
	sample->imageSample = Point2(image[samplePos]);
	sample->lensSample = Point2(lens[samplePos]);
	sample->time = time[samplePos];

	// Generate samples for integrator
	for(int j = 0; j < sample->n1D.size(); j++)
			LatinHypercube1D(sample->oneD[j], sample->n1D[j], rng);
	for(int j = 0; j < sample->n2D.size(); j++)
			LatinHypercube2D(sample->twoD[j], sample->n2D[j], rng);

	samplePos++;

	return 1;
}

